from lib.tencent.hunyuan_lite import HunyuanLite
from lib.baidu.ernie_speed import ErnieSpeed


class Question:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "system": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "给你一个任务，对输入内容进行改写，改写成描述一副画的内容，改写的内容只能是对输入内容中已有的实物元素进行描述，或者添加一些简单的实物元素进行描述，改写的内容要符合实际且要精准描述，改写的内容画面要突出输入的内容，改写内容不要有情感、气氛、比喻，假想，联想等非实物的内容，改写内容不要漏掉输入内容中的实体元素，改写内容不要有多余的解释，然后返回改写的内容给我，",
                    },
                ),
                "user": ("STRING", {"multiline": True, "default": ""}),
                "engine": (
                    ["ernie_speed_128k", "hunyuan_lite"],
                    {"default": "ernie_speed_128k"},
                ),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("result",)

    FUNCTION = "handle"

    OUTPUT_NODE = False

    CATEGORY = "fcl/question"

    history_dict = {}

    def handle(self, system="", user="", engine=""):

        messages = []

        if engine == "hunyuan_lite":
            if system != "":
                messages.append({"Role": "system", "Content": system})
            if user != "":
                messages.append({"Role": "user", "Content": user})
            print(messages)
            hyl = HunyuanLite()
            ret_msg = hyl.getMessage(hyl.chat(messages=messages))
        else:
            if system != "":
                messages.append({"role": "user", "content": system})
                messages.append({'role': 'assistant', 'content': '好的，请输入你的内容。'})
            if user != "":
                messages.append({"role": "user", "content": user})
            print(messages)
            es = ErnieSpeed()
            ret_msg = es.getMessage(es.chat(messages=messages))

        return (ret_msg,)


NODE_CLASS_MAPPINGS = {"FCL_QUESTION": Question}

NODE_DISPLAY_NAME_MAPPINGS = {"FCL_QUESTION": "fcl dev question"}
